package com.atao.baokede.controller;

import com.atao.baokede.chatMemory.MysqlChatMemory;
import com.atao.baokede.pojo.AiDialogue;
import com.atao.baokede.pojo.Result;
import com.atao.baokede.pojo.User;
import com.atao.baokede.service.AiService;
import com.atao.baokede.service.UserService;
import com.atao.baokede.utils.ThreadLocalUtil;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Lazy;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;

/**
 * AI层
 */
@RestController
@RequestMapping("/ai")
public class AiController {
    @Autowired
    private ChatClient chatClient;

    @Autowired
    private MysqlChatMemory mysqlChatMemory;

    @Autowired
    private VectorStore vectorStore;

    @Autowired
    private ToolCallback[] tools;

    @Autowired
    private UserService userService;

    @Autowired
    @Lazy
    private ToolCallbackProvider toolCallbackProvider;

    @Autowired
    private AiService aiService;

    /**
     * AI对话
     * @param prompt 用户的提示词
     * @return AI回答
     */
    @GetMapping(value = "/streamChat", produces = "text/event-stream;charset=UTF-8")
    public Flux<String> stream(@RequestParam String prompt) {
        // 从ThreadLocal获取用户信息
        Map<String, Object> map = ThreadLocalUtil.get();
        Long userId = ((Number) map.get("id")).longValue();
        String conversationId = userId.toString();

        User user = userService.getUserInfoById(userId);
        var historyMessages = mysqlChatMemory.get(conversationId);
        var allMessages = new java.util.ArrayList<>(historyMessages);

        allMessages.add(new org.springframework.ai.chat.messages.UserMessage(prompt));

        StringBuilder completeResponse = new StringBuilder();

        return chatClient.prompt()
                .system(promptSystemSpec -> promptSystemSpec.param("userName", user.getUsername()))
                .messages(allMessages)
                .options(DashScopeChatOptions.builder()
                        .withMaxToken(512)
                        .build())
                .advisors(new QuestionAnswerAdvisor(vectorStore))
                .toolCallbacks(tools)
                .toolCallbacks(toolCallbackProvider)
                .stream()
                .content()
                .map(chunk -> {
                    // 1. 将当前数据块追加到completeResponse（原始数据，包含\n）
                    completeResponse.append(chunk);
                    // 2. 对此发往前端的数据块进行处理：将换行符\n替换为<br>
                    String chunkForFrontend = chunk.replace("\n", "<br>");
                    // 3. 确保返回的数据符合SSE格式 ("data: ...\n\n")
                    // 如果你的chatClient返回的内容已经是纯内容而非SSE格式，你需要包装它
                    return chunkForFrontend + "\n\n";
                })
                .doOnComplete(() -> {
                    String fullResponse = completeResponse.toString(); // 这是原始的、包含\n的完整响应
                    System.out.println(fullResponse);
                    if (!fullResponse.isEmpty()) {
                        // 将原始响应（含\n）保存到数据库
                        mysqlChatMemory.saveDialogue(conversationId, prompt, fullResponse);
                    }
                    System.out.println("Stream completed for user: " + userId + ", response length: " + fullResponse.length());
                })
                .doOnError(error -> {
                    System.err.println("Stream error for user " + userId + ": " + error.getMessage());
                })
                .concatWith(Flux.defer(() -> {
                    if (completeResponse.length() >= 512) {
                        // 截断信息也需要转换为SSE格式并处理换行（如果需要）
                        String truncationMessage = "{\"status\":\"truncated\",\"message\":\"The response was truncated due to maxTokens limit.\"}";
                        String escapedTruncationMessage = truncationMessage.replace("\n", "<br>");
                        return Flux.just(escapedTruncationMessage + "\n\n");
                    }
                    return Flux.empty();
                }))
                .concatWith(Flux.just("[END]\n\n")); // 结束标记也转换为SSE格式
    }


    /**
     * 删除当前用户与AI的所有聊天记录
     */
    @GetMapping("/deleteAllChat")
    public Result deleteAllChat() {
        aiService.deleteAllChat();
        return Result.success();
    }

    /**
     * 获取当前用户与AI的所有聊天记录
     */
    @GetMapping("/getAllChat")
    public Result<List<AiDialogue>> getAllChat() {
        return Result.success(aiService.getAllChat());
    }
}